﻿#include "rtmppusher.h"

// 大小端字节序转换
#define HTON16(x) ( (x >> 8 & 0x00FF) | (x << 8 & 0xFF00) )
#define HTON24(x) ( (x >> 16 & 0x0000FF) | (x & 0x00FF00) | (x << 16 & 0xFF0000) )
#define HTON32(x) ( (x >> 24 & 0x000000FF) | (x >> 8 & 0x0000FF00) | (x << 8 & 0x00FF0000) | (x << 24 & 0xFF000000) )
#define HTONTIME(x) ( (x >> 16 & 0x000000FF) | (x & 0x0000FF00) | (x << 16 & 0x00FF0000) | (x & 0xFF000000) )

int64_t hexArrayToDec(char *array, int len)
{
    if(array == nullptr)
        return -1;

    int64_t result = 0;
    for(int i=0; i<len; i++)
        result = result * 256 + (unsigned char)array[i];

    return result;
}

// 从文件读取指定字节
bool ReadFP(char* pBuf, int nSize, FILE* pFile)
{
    return (fread(pBuf, 1, nSize, pFile) == nSize);
}

// 从文件读取1个字节整数
bool ReadU8(uint8_t* u8, FILE* fp)
{
    return ReadFP((char*)u8, 1, fp);
}

// 从文件读取2个字节整数
bool ReadU16(uint16_t* u16, FILE* fp)
{
    if (!ReadFP((char*)u16, 2, fp))
        return false;

    *u16 = HTON16(*u16);
    return true;
}

// 从文件读取3个字节整数
bool ReadU24(uint32_t* u24, FILE* fp)
{
    if (!ReadFP((char*)u24, 3, fp))
        return false;

    *u24 = HTON24(*u24);
    return true;
}

// 从文件读取4个字节整数
bool ReadU32(uint32_t* u32, FILE* fp)
{
    if (!ReadFP((char*)u32, 4, fp))
        return false;

    *u32 = HTON32(*u32);
    return true;
}

// 从文件读取4个字节时间戳
bool ReadTime(uint32_t* utime, FILE* fp)
{
    if (!ReadFP((char*)utime, 4, fp))
        return false;

    *utime = HTONTIME(*utime);
    return true;
}

// 从文件预读1个字节整数
bool PeekU8(uint8_t* u8, FILE* fp)
{
    if (!ReadFP((char*)u8, 1, fp))
        return false;

    fseek(fp, -1, SEEK_CUR);
    return true;
}

char * put_byte(char *output, uint8_t nVal)
{
    output[0] = nVal;
    return output + 1;
}
char * put_be16(char *output, uint16_t nVal)
{
    output[1] = nVal & 0xff;
    output[0] = nVal >> 8;
    return output + 2;
}
char * put_be24(char *output, uint32_t nVal)
{
    output[2] = nVal & 0xff;
    output[1] = nVal >> 8;
    output[0] = nVal >> 16;
    return output + 3;
}
char * put_be32(char *output, uint32_t nVal)
{
    output[3] = nVal & 0xff;
    output[2] = nVal >> 8;
    output[1] = nVal >> 16;
    output[0] = nVal >> 24;
    return output + 4;
}
char *  put_be64(char *output, uint64_t nVal)
{
    output = put_be32(output, nVal >> 32);
    output = put_be32(output, nVal);
    return output;
}
char * put_amf_string(char *c, const char *str)
{
    uint16_t len = strlen(str);
    c = put_be16(c, len);
    memcpy(c, str, len);
    return c + len;
}
char * put_amf_double(char *c, double d)
{
    *c++ = AMF_NUMBER;  /* type: Number */
    {
        unsigned char *ci, *co;
        ci = (unsigned char *)&d;
        co = (unsigned char *)c;
        co[0] = ci[7];
        co[1] = ci[6];
        co[2] = ci[5];
        co[3] = ci[4];
        co[4] = ci[3];
        co[5] = ci[2];
        co[6] = ci[1];
        co[7] = ci[0];
    }
    return c + 8;
}


/*********** RtmpPusher ***********/
RtmpPusher::RtmpPusher()
{
    isRun = false;

    aacSeqHeader = nullptr;
    avcSeqHeader = nullptr;
    metaMsg = nullptr;
}

RtmpPusher::~RtmpPusher()
{
    release();

    if(this->avcSeqHeader)
        delete[] this->avcSeqHeader;
    if(this->aacSeqHeader)
        delete[] this->aacSeqHeader;

    quit();
    wait();
}

void RtmpPusher::setRunState(bool start)
{
    isRun = start;

    if(start)
        this->start();
}

void RtmpPusher::testPushFlvFile()
{
    //给RTMP申请内存
    RTMP *rtmp = RTMP_Alloc();
    if (!rtmp)
    {
        qDebug() << "alloc rtmp 失败";
        return;
    }

    //初始化
    RTMP_Init(rtmp);

    //设置推流url
    char *url = "rtmp://127.0.0.1/live/test";
    int ret = RTMP_SetupURL(rtmp, url);
    if (!ret)
    {
        qDebug() << "设置RTMP推流地址失败" << url;
        return;
    }

    //设置超时时间
    rtmp->Link.timeout = 5;

    //启用RTMP写入功能
    RTMP_EnableWrite(rtmp);

    //链接服务器
    ret = RTMP_Connect(rtmp, 0);
    if (!ret)
    {
        qDebug() << "链接服务器失败" << url;
        return;
    }

    //链接流
    ret = RTMP_ConnectStream(rtmp, 0);
    if (!ret)
    {
        qDebug() << "链接流失败" << url;
        return;
    }

    //本地文件
    std::string path = filePath.toStdString();
    FILE* pFile = fopen(path.data(), "rb");

    //跳过FLV文件头
    fseek(pFile, 9, SEEK_SET);

    //初使化RTMP报文
    RTMPPacket packet;
    RTMPPacket_Reset(&packet);
    packet.m_body = NULL;
    packet.m_chunk = NULL;
    packet.m_nInfoField2 = rtmp->m_stream_id;

    //记录一个推流的开始时间
    uint32_t startTime = av_gettime() / 1000;

    qDebug() << "push file start";
    while(isRun)
    {
        //跳过PreTag
        uint32_t pretagsize = 0;
        if (!ReadU32(&pretagsize, pFile))
            break;

        //读取TAG头
        uint8_t type = 0;
        if (!ReadU8(&type, pFile))
            break;

        uint32_t datalen = 0;
        if (!ReadU24(&datalen, pFile))
            break;

        uint32_t timestamp = 0;
        if (!ReadTime(&timestamp, pFile))
            break;

        uint32_t streamid = 0;
        if (!ReadU24(&streamid, pFile))
            break;

        RTMPPacket_Alloc(&packet, datalen);
        if (fread(packet.m_body, 1, datalen, pFile) != datalen)
            break;

        //组织报文并发送
        packet.m_headerType = RTMP_PACKET_SIZE_LARGE;
        packet.m_packetType = type;
        packet.m_hasAbsTimestamp = 0;
        packet.m_nChannel = 6;
        packet.m_nTimeStamp = timestamp;
        packet.m_nBodySize = datalen;

        if (!RTMP_SendPacket(rtmp, &packet, 0))
        {
            qDebug() << "Send Error!";
            break;
        }

        //延时均为毫秒级 进行频率控制
        uint32_t elapsed = (av_gettime() / 1000 - startTime);
        uint32_t offset = timestamp - elapsed;
        if (offset > 5 && offset < 1000)
            QThread::msleep(offset);

        qDebug() << "type" << type << "timestamp" << timestamp << "size" << datalen;

        RTMPPacket_Free(&packet);
    }

    qDebug() << "push file finish";

    //释放
    if (rtmp)
    {
        RTMP_Close(rtmp);
        RTMP_Free(rtmp);
    }
}

void RtmpPusher::initSender()
{
    release();
}

void RtmpPusher::release()
{
    QMutexLocker locker(&mutex);

    for(int i=packetList.size() - 1; i >= 0; i--)
    {
        ReadyPacket *ready = packetList.at(i);
        AVPacket *packet = ready->packet;
        av_packet_free(&packet);

        delete ready;
        ready = nullptr;
    }

    packetList.clear();
}

void RtmpPusher::writePacket(ReadyPacket *packet)
{
    QMutexLocker locker(&mutex);
    packetList.append(packet);
}

void RtmpPusher::setMetaMsg(FLVMetadataMsg *metaMsg)
{
    if(this->metaMsg)
        delete this->metaMsg;

    this->metaMsg = metaMsg;
}

void RtmpPusher::setAudioSeqHeader(char *aacSeqHeader, int aacSeqHeaderLen)
{
    //通过AVCodecContent中的extradata设置
    if(hasAudio)
    {
        if(this->aacSeqHeader)
            delete[] this->aacSeqHeader;

        this->aacSeqHeaderLen = aacSeqHeaderLen;
        this->aacSeqHeader = new char[aacSeqHeaderLen];
        memcpy(this->aacSeqHeader, aacSeqHeader, aacSeqHeaderLen);
    }
}

void RtmpPusher::setVideoSeqHeader(char *avcSeqHeader, int avcSeqHeaderLen)
{
    //通过AVCodecContent中的extradata设置
    if(hasVideo)
    {
        if(this->avcSeqHeader)
            delete[] this->avcSeqHeader;

        this->avcSeqHeaderLen = avcSeqHeaderLen;
        this->avcSeqHeader = new char[avcSeqHeaderLen];
        memcpy(this->avcSeqHeader, avcSeqHeader, avcSeqHeaderLen);
    }
}

void RtmpPusher::parseVideoAVPacket(char *packetData, int packetSize, uint32_t ptsTime, uint32_t dtsTime)
{
    if(packetData == nullptr || packetSize < 4)
        return;

    //已解析长度
    int hasParseLen = 0;

    //解封装后的AVpacket中存在多帧组合的情况如SPS、PPS、IDR 组包RTP需分开每一帧；部分也有缺少SPS、PPS，需填充
    while(1)
    {
        //解封装后的AVpacket数据格式：四字节大小+实际数据
        char *curPacketData = packetData + hasParseLen;
        int curLen = hexArrayToDec(curPacketData, 4);
        hasParseLen += curLen + 4;
        char *curFrame = curPacketData + 4;

        //RTMP推流开始发送了音、视频解码信息，因此这里只发送实际音视频数据帧即可
        uint8_t typeValue = curFrame[0] & 0x1F;

        //得到单独一帧组包发送 I帧 P帧
        if(typeValue == 5 || typeValue == 1)
            sendH264RawData(curFrame, curLen, ptsTime, dtsTime);

        if(hasParseLen == packetSize)
            break;
    }
}

void RtmpPusher::parseAudioAVPacket(char *packetData, int packetSize, uint32_t ptsTime)
{
    if(packetData == nullptr || packetSize < 4)
        return;

    //音频AVPacket一般是单独一帧压缩数据 不需额外分帧处理
    sendAACRawData(packetData, packetSize, ptsTime);
}

bool RtmpPusher::sendMetadata(FLVMetadataMsg *metadata)
{
    if (metadata == NULL)
        return false;

    char body[1024] = {0};

    char *p = (char *)body;   
    p = put_byte(p, AMF_STRING);
    p = put_amf_string(p, "@setDataFrame");

    p = put_byte(p, AMF_STRING);
    p = put_amf_string(p, "onMetaData");

    p = put_byte(p, AMF_OBJECT);
    if(metadata->has_video)
    {
        p = put_amf_string(p, "width");
        p = put_amf_double(p, metadata->width);

        p = put_amf_string(p, "height");
        p = put_amf_double(p, metadata->height);

        p = put_amf_string(p, "framerate");
        p = put_amf_double(p, metadata->framerate);

        p = put_amf_string(p, "videodatarate");
        p = put_amf_double(p, metadata->videodatarate);

        p = put_amf_string(p, "videocodecid");
        p = put_amf_double(p, FLV_CODECID_H264);
    }

    if(metadata->has_audio)
    {
        p = put_amf_string(p, "audiodatarate");
        p = put_amf_double(p, (double)metadata->audiodatarate);

        p = put_amf_string(p, "audiosamplerate");
        p = put_amf_double(p, (double)metadata->audiosamplerate);

        p = put_amf_string(p, "audiosamplesize");
        p = put_amf_double(p, (double)metadata->audiosamplesize);

        p = put_amf_string(p, "stereo");
        p = put_byte(p, AMF_BOOLEAN);
        p = put_byte(p, (metadata->channles == 2));

        p = put_amf_string(p, "audiocodecid");
        p = put_amf_double(p, (double)FLV_CODECID_AAC);
    }

    p = put_amf_string(p, "");
    p = put_byte(p, AMF_OBJECT_END);

    //此时设为起始时间更准确
    pushStartTime = av_gettime() / 1000;
    return sendRtmpPacket(RTMP_PACKET_TYPE_INFO, (uint8_t*)body, p-body, 0);
}

bool RtmpPusher::sendAACSpecificConfig()
{
    uint8_t aac_spec_[4] = {0};
    aac_spec_[0] = 0xAF;     // 一般默认 以下方sequence header为准
    aac_spec_[1] = 0x00;     // 0:aac sequence header  1:aac rawdata

    //SpecificConfig可以由profile samplerate channel生成  这里直接使用解码器上下文中的extradata
    if(hasAudio)
    {
        aac_spec_[2] = aacSeqHeader[0];
        aac_spec_[3] = aacSeqHeader[1];
    }

    return sendRtmpPacket(RTMP_PACKET_TYPE_AUDIO, aac_spec_, 4, 0);
}

bool RtmpPusher::sendH264SequenceHeader()
{
    uint8_t body[512] = {0};
    body[0] = 0x17; // 1:keyframe  7:AVC
    body[1] = 0x00; // 0:AVC sequence header 1:vc rawdata

    body[2] = 0x00;
    body[3] = 0x00;
    body[4] = 0x00; //timestamp 0

    //AVCDecoderConfigurationRecord 可以由视频长、宽、帧率等中信息生成  这里直接使用解码器上下文中的extradata
    memcpy(body + 5, avcSeqHeader, avcSeqHeaderLen);
    int size = avcSeqHeaderLen + 5;

    return sendRtmpPacket(RTMP_PACKET_TYPE_VIDEO, (uint8_t *)body, size, 0);
}

bool RtmpPusher::sendAACRawData(char *rawData, int size, uint32_t ptsTime)
{
    int totalSize = size + 2;
    char *body = new char[totalSize];
    body[0] = 0xAF;  // 一般默认 以下方sequence header为准
    body[1] = 0x01;  // 0:aac sequence header  1:aac rawdata

    memcpy(body + 2, rawData, size);

    //完整rtmp数据交由librtmp发送
    int ret = sendRtmpPacket(RTMP_PACKET_TYPE_AUDIO, (uint8_t *)body, totalSize, ptsTime);
    if(ret != 1)
        qDebug() << "aac sendRtmpPacket error";

    delete[] body;
    body = nullptr;
    return ret;
}

bool RtmpPusher::sendH264RawData(char *rawData, int size, uint32_t ptsTime, uint32_t dtsTime)
{
    int totalSize = size + 9;
    char *body = new char[totalSize];

    //帧类型判断 I帧 P帧
    uint8_t typeValue = rawData[0] & 0x1F;
    if(typeValue == 5)
        body[0] = 0x17; // 1:keyframe  7:AVC
    else
        body[0] = 0x27; // 2:interframe  7:AVC

    body[1] = 0x01; // 0:AVC sequence header 1: rawdata

    //pts >= dts offset恒大于0
    uint32_t offset = ptsTime - dtsTime;
    body[2] = offset & 0xFF0000;
    body[3] = offset & 0x00FF00;
    body[4] = offset & 0x0000FF;

    // NALU size
    body[5] = size >> 24;
    body[6] = size >> 16;
    body[7] = size >> 8;
    body[8] = size & 0xff;
    memcpy(body + 9, rawData, size);

    //完整rtmp数据交由librtmp发送
    int ret = sendRtmpPacket(RTMP_PACKET_TYPE_VIDEO, (uint8_t *)body, totalSize, ptsTime);
    if(ret != 1)
        qDebug() << "h264 sendRtmpPacket error";

    delete[] body;
    body = nullptr;
    return ret;
}

int RtmpPusher::sendRtmpPacket(int packetType, uint8_t *data, int size, uint32_t timeStamp)
{
    RTMPPacket packet;
    RTMPPacket_Reset(&packet);
    RTMPPacket_Alloc(&packet, size);

    packet.m_packetType = packetType;
    if(packetType == RTMP_PACKET_TYPE_AUDIO)
        packet.m_nChannel = RTMP_AUDIO_CHANNEL;
    else if(packetType == RTMP_PACKET_TYPE_VIDEO)
        packet.m_nChannel = RTMP_VIDEO_CHANNEL;
    else
        packet.m_nChannel = RTMP_NETWORK_CHANNEL;

    packet.m_headerType = RTMP_PACKET_SIZE_LARGE;
    packet.m_nTimeStamp = timeStamp;
    packet.m_nInfoField2 = rtmp->m_stream_id;
    packet.m_nBodySize = size;
    memcpy(packet.m_body, data, size);

    //延时均为毫秒级 进行频率控制
    uint32_t elapsed = (av_gettime() / 1000 - pushStartTime);
    uint32_t offset = timeStamp - elapsed;
    if (offset > 10 && offset < 1000)
        QThread::msleep(offset);

    int ret = RTMP_SendPacket(rtmp, &packet, 0);
    RTMPPacket_Free(&packet);
    return ret;
}


void RtmpPusher::run()
{
#if 0
    //推送本地flv文件
    testPushFlvFile();
#else
    //给RTMP申请内存
    rtmp = RTMP_Alloc();
    if (!rtmp)
    {
        qDebug() << "alloc rtmp 失败";
        return;
    }

    //初始化
    RTMP_Init(rtmp);

    //设置推流url
    char *url = "rtmp://192.168.10.6/live/test";
    int ret = RTMP_SetupURL(rtmp, url);
    if (!ret)
    {
        qDebug() << "设置RTMP推流地址失败" << url;
        return;
    }

    //设置超时时间
    rtmp->Link.timeout = 5;

    //启用RTMP写入功能
    RTMP_EnableWrite(rtmp);

    //链接服务器
    ret = RTMP_Connect(rtmp, 0);
    if (!ret)
    {
        qDebug() << "链接服务器失败" << url;
        return;
    }

    //链接流
    ret = RTMP_ConnectStream(rtmp, 0);
    if (!ret)
    {
        qDebug() << "链接流失败" << url;
        return;
    }

    qDebug() << "开始推流";
    release();

    //metaData作为推流的开始时间
    sendMetadata(this->metaMsg);

    if(hasAudio)
        sendAACSpecificConfig();
    if(hasVideo)
        sendH264SequenceHeader();

    while (isRun)
    {
        //读取队列当前包 只读取减少锁粒度
        QList<ReadyPacket *> list;
        int size = packetList.size();
        mutex.lock();
        for(int i=0; i<size; i++)
        {
            ReadyPacket *ready = packetList.first();
            list.append(ready);
            packetList.pop_front();
        }
        mutex.unlock();

        for(int i=0; i<list.size(); i++)
        {
            ReadyPacket *ready = list.at(i);
            AVPacket *packet = ready->packet;
            int type = ready->type;

            if(type == TypeVideo)
                parseVideoAVPacket((char *)packet->data, packet->size, ready->ptsTime, ready->dtsTime);
            else if(type == TypeAudio)
                parseAudioAVPacket((char *)packet->data, packet->size, ready->ptsTime);

            av_packet_free(&packet);
            delete ready;
            ready = nullptr;
        }

        QThread::msleep(1);
    }

    qDebug() << "结束推流" << endl;

    //释放
    if (rtmp)
    {
        RTMP_Close(rtmp);
        RTMP_Free(rtmp);
    }

#endif
}
